from torch import nn

class ANN(nn.Module):
    def __init__(
            self, width, height, 
            depth=1, class_num=10, hidden=1024
        ) -> None:
        super().__init__()
        self.size = width*height
        self.class_num = class_num
        self.hidden_dim = hidden
        
        self.flat = nn.Flatten()
        self.linear = nn.Sequential()
        in_dim = self.size
        out_dim = class_num
        for i in range(depth):
            if i < depth-1:
                out_dim = hidden
            self.linear.append(
                nn.Linear(in_dim, out_dim)
            )
            in_dim = hidden
            if i == depth-2:
                out_dim = class_num
    def forward(self,x):
        x = self.flat(x)
        x = self.linear(x)
        return x

class CNN(nn.Module):
    def __init__(self,class_num) -> None:
        super().__init__()
        self.LeNet = nn.Sequential(
            nn.Conv2d(1,6,5),
            nn.MaxPool2d(2),
            nn.Conv2d(6,16,5),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(16*4*4,128),
            nn.ReLU(),
            nn.Linear(128,32),
        )
        self.out = nn.Sequential(
            nn.ReLU(),
            nn.Linear(32,class_num)
        )

    def forward(self,x):
        x = self.LeNet(x)
        return self.out(x)